// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// #include "postprocess_op.h"
// #include "preprocess_op.h"
// #include "base_util/utils.h"
#include "utils/processor.h"

#include "utils/transforms.h"
#include "utils/model_config.h"

#include "ocr_rec.h"


namespace paddle_infer {

OCRRecognizer::OCRRecognizer(LogInfo *lg):
  log_ifo(lg) {}

int OCRRecognizer::Run(cv::Mat& img, 
                      std::vector<BaseInfoPD*> &rec_ifo, 
                      PDModelManager *mng,
                      int md_idx ) 
{
  if (mng->configs[md_idx]->algorithm == pd_text_rec){
    return RunCRNN(img, rec_ifo, mng, md_idx);
  }else {
    return pd_not_support_algorithm;
  }
  return pd_not_support_algorithm;
}


int OCRRecognizer::RunCRNN(cv::Mat& img, 
                          std::vector<BaseInfoPD*> &rec_ifo, 
                          PDModelManager *mng,
                          int md_idx) 
{
  int ret =0;
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.0 RunCRNN");

  if (img.empty()){return PD_IMAGE_EMPTY_ERROR;}
  auto& cfg = mng->configs[md_idx];
  ImageBlob img_blob = ImageBlob();
  cfg->transforms->run(img, img_blob, cfg);
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.1 RunCRNN");

  auto input_names = mng->predictors[md_idx]->GetInputNames();
  auto in_tensor = mng->predictors[md_idx]->GetInputTensor(input_names[0]);
  in_tensor->Reshape({1, cfg->channels, img_blob.new_im_shape[0], img_blob.new_im_shape[1]});
  in_tensor->CopyFromCpu(img_blob.im_data.data());

  mng->predictors[md_idx]->ZeroCopyRun();
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.2 RunCRNN");

  auto output_names = mng->predictors[md_idx]->GetOutputNames();
  if (log_ifo->log_level_4) { log_tensor_shape("output_names", output_names); }
  
  auto out_tensor = mng->predictors[md_idx]->GetOutputTensor(output_names[0]);
  std::vector<int> output_shape = out_tensor->shape();
  if (log_ifo->log_level_4) { log_tensor_shape("output_tensor", output_shape); }
  
  int size = std::accumulate(output_shape.begin(), output_shape.end(), 1, std::multiplies<int>());
  std::vector<float> output_data(size);
  out_tensor->CopyToCpu(output_data.data());
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.3 RunCRNN");

  // ctc decode
  std::string text = ""; 
  int argmax_idx;
  int last_index = 0;
  float score = 0.f;
  int count = 0;
  float max_value = 0.0f;

  for (int n = 0; n < output_shape[1]; n++) {
    argmax_idx =
        int(PostProcessor::argmax(&output_data[n * output_shape[2]],
                            &output_data[(n + 1) * output_shape[2]]));
    max_value =
        float(*std::max_element(&output_data[n * output_shape[2]],
                                &output_data[(n + 1) * output_shape[2]]));
      if (argmax_idx > 0 && (!(n > 0 && argmax_idx == last_index)) 
          && argmax_idx < mng->configs[md_idx]->label_list.size()) {
        score += max_value;
        count += 1;
        text += mng->configs[md_idx]->label_list[argmax_idx];
      }
      last_index = argmax_idx;
      // if (argmax_idx > 0  && argmax_idx < mng->configs[model_idx]->label_list.size()) {       // && argmax_idx != last_index
      //   score += max_value;
      //   count += 1;
      //   text += mng->configs[model_idx]->label_list[argmax_idx];
      // }
  }
  score /= count;

  if (log_ifo->log_level_4){
    spdlog::get("logger")->info("::RunCRNN --> text: {} \tscore:{}",  text, score);
  }
  rec_ifo.push_back(new TextRecInfoPD(score, text));
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.4 RunCRNN");

  return 0;
}


} // namespace paddle_infer
